
# parses a bunch of files and creates histograms

from StringIO import StringIO
import struct

import csv
import os
import re
import sys

import math
import numpy
import matplotlib
import matplotlib.figure
from matplotlib.backends.backend_agg import FigureCanvasAgg

from collections import defaultdict

#from mpl_toolkits.mplot3d import Axes3D

# assumes the following naming convention:
# YYYY-MM-DD_[gene]_[drug]_[dose]_t[time].fcs
f_re = re.compile(r'\d{4}-\d{2}-\d{2}_(\w+)_(\w+)_([0-9\.]+)_t(\d+)',
                  flags=re.I)

def fcsextract(filename):
    """
    Attempts to parse an FCS (flow cytometry standard) file

    Parameters: filename
        filename: path to the FCS file

    Returns: (vars,events)
    	vars: a dictionary with the KEY/VALUE pairs found in the HEADER
    	this includes the standard '$ABC' style FCS variable as well as any 
    	custom variables added to the header by the machine or operator
	
    	events: an [N x D] matrix of the data (as a Python list of lists)
    	i.e. events[99][2] would be the value at the 3rd dimension
    	of the 100th event
    """
    fcs_file_name = filename

    fcs = open(fcs_file_name,'rb')
    header = fcs.read(58)
    version = header[0:6].strip()
    text_start = int(header[10:18].strip())
    text_end = int(header[18:26].strip())
    data_start = int(header[26:34].strip())
    data_end = int(header[34:42].strip())
    analysis_start = int(header[42:50].strip())
    analysis_end = int(header[50:58].strip())

    #print "Parsing TEXT segment"
    # read TEXT portion
    fcs.seek(text_start)
    delimeter = fcs.read(1)
    # First byte of the text portion defines the delimeter
    #print "delimeter:",delimeter
    text = fcs.read(text_end-text_start+1)

    #Variables in TEXT poriton are stored "key/value/key/value/key/value"
    keyvalarray = text.split(delimeter)
    fcs_vars = {}
    fcs_var_list = []
    # Iterate over every 2 consecutive elements of the array
    for k,v in zip(keyvalarray[::2],keyvalarray[1::2]):
        fcs_vars[k] = v
        fcs_var_list.append((k,v)) # Keep a list around so we can print them in order

    #from pprint import pprint; pprint(fcs_var_list)
    if data_start == 0 and data_end == 0:
        data_start = int(fcs_vars['$DATASTART'])
        data_end = int(fcs_vars['$DATAEND'])

    num_dims = int(fcs_vars['$PAR'])
    #print "Number of dimensions:",num_dims

    num_events = int(fcs_vars['$TOT'])
    #print "Number of events:",num_events

    # Read DATA portion
    fcs.seek(data_start)
    #print "# of Data bytes",data_end-data_start+1
    data = fcs.read(data_end-data_start+1)

    # Determine data format
    datatype = fcs_vars['$DATATYPE']
    if datatype == 'F':
        datatype = 'f' # set proper data mode for struct module
        #print "Data stored as single-precision (32-bit) floating point numbers"
    elif datatype == 'D':
        datatype = 'd' # set proper data mode for struct module
        #print "Data stored as double-precision (64-bit) floating point numbers"
    else:
        assert False,"Error: Unrecognized $DATATYPE '%s'" % datatype
    
    # Determine endianess
    endian = fcs_vars['$BYTEORD']
    if endian == "4,3,2,1":
        endian = ">" # set proper data mode for struct module
        #print "Big endian data format"
    elif endian == "1,2,3,4":
        #print "Little endian data format"
        endian = "<" # set proper data mode for struct module
    else:
        assert False,"Error: This script can only read data encoded with $BYTEORD = 1,2,3,4 or 4,3,2,1"

    # Put data in StringIO so we can read bytes like a file    
    data = StringIO(data)

    #print "Parsing DATA segment"
    # Create format string based on endianeness and the specified data type
    format = endian + str(num_dims) + datatype
    datasize = struct.calcsize(format)
    #print "Data format:",format
    #print "Data size:",datasize
    events = []
    # Read and unpack all the events from the data
    for e in range(num_events):
        event = struct.unpack(format,data.read(datasize))
        events.append(event)
    
    fcs.close()
    return fcs_vars,events
    

def filter_data(data):
    # insert some filtering stuff here...
    return [(row['FSC-A'],row['SSC-A'],row['FITC-A']) for row in data]

def histo(f_list):
    time_courses = defaultdict(list)

    for f in f_list:
        fcs_vars,events = fcsextract(f)
        f_name = os.path.basename(f)
        gene,drug,dose,time = f_re.match(f_name).groups()
        
        header = [fcs_vars['$P%dN' % (i+1)] for i in range(len(events[0]))]
        
        data = filter_data([dict(zip(header,e)) for e in events])
        
        d = [math.log(d[2],10) for d in data if d[2] > 0.0]
        #print min(d), max(d)
        
        time_courses[gene].append((int(time), d))

        
    for gene in time_courses:
        time_courses[gene].sort()

#     gene_medians = dict()
# 
#     for gene in time_courses:
#         total_data = sum(time_courses[gene][t][1] for t in range(11), [])
#         gene_medians[gene] = sorted(total_data)[len(total_data)/2]

    for t in range(11):
        fig = matplotlib.figure.Figure()

        for i,gene in enumerate(sorted(time_courses.keys())):
            ax = fig.add_subplot(5,6,i+1)
            ax.hist(time_courses[gene][t][1], 100)
            ax.text(0.05, 0.95, gene,
                    horizontalalignment='left',
                    verticalalignment='top',
                    transform=ax.transAxes)
            ax.get_yaxis().set_visible(False)
            ax.get_xaxis().set_visible(False)
            ax.set_xlim(xmin=0.0, xmax=6.0)
    
        p = r'/Users/ucsf/Desktop/team_jkrw/cytometry/cerulenin_multi_hist/%s_multi.svg'
        FigureCanvasAgg(fig).print_figure(p % t)

if __name__ == '__main__':
    histo(sys.argv[1:])
